import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np

# 1. 配置设备（优先用 GPU）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 2. 加载 MNIST 数据集 & 预处理
transform = transforms.Compose([
    transforms.ToTensor(),  # 转 Tensor，归一化到 [0,1]
    transforms.Normalize((0.5,), (0.5,))  # 标准化：(x-0.5)/0.5 → 范围 [-1,1]
])

# 训练集
train_dataset = torchvision.datasets.MNIST(
    root='./data', train=True, download=True, transform=transform
)
# 测试集
test_dataset = torchvision.datasets.MNIST(
    root='./data', train=False, download=True, transform=transform
)

# 数据加载器（batch 训练）
batch_size = 64
train_loader = torch.utils.data.DataLoader(
    train_dataset, batch_size=batch_size, shuffle=True
)
test_loader = torch.utils.data.DataLoader(
    test_dataset, batch_size=batch_size, shuffle=False
)


# 3. 构建两层神经网络
class TwoLayerNet(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(TwoLayerNet, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)  # 第一层：输入层→隐藏层
        self.relu = nn.ReLU()  # 激活函数
        self.fc2 = nn.Linear(hidden_size, num_classes)  # 第二层：隐藏层→输出层

    def forward(self, x):
        # 展平输入（MNIST 是 28×28，转成 784 维向量）
        x = x.view(-1, 28 * 28)
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        return out


# 初始化模型（输入 784，隐藏层维度可调，输出 10 分类）
input_size = 28 * 28
hidden_size = 128  # 可调整，影响精度和训练速度
num_classes = 10
model = TwoLayerNet(input_size, hidden_size, num_classes).to(device)

# 4. 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()  # 多分类交叉熵
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)  # Adam 优化器

# 5. 训练模型 + 记录指标
num_epochs = 10  # 可调整，建议先小试，再增大
train_losses = []
train_accuracies = []
test_losses = []
test_accuracies = []

for epoch in range(num_epochs):
    # --- 训练阶段 ---
    model.train()
    train_correct = 0
    train_total = 0
    train_running_loss = 0.0

    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)

        # 前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)

        # 反向传播 + 优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # 统计损失和正确数
        train_running_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        train_total += labels.size(0)
        train_correct += (predicted == labels).sum().item()

    # 计算 epoch 级训练指标
    train_epoch_loss = train_running_loss / len(train_loader)
    train_epoch_accuracy = 100 * train_correct / train_total
    train_losses.append(train_epoch_loss)
    train_accuracies.append(train_epoch_accuracy)

    # --- 测试阶段 ---
    model.eval()
    test_correct = 0
    test_total = 0
    test_running_loss = 0.0

    with torch.no_grad():  # 测试时无需梯度
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)

            outputs = model(images)
            loss = criterion(outputs, labels)

            test_running_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            test_total += labels.size(0)
            test_correct += (predicted == labels).sum().item()

    # 计算 epoch 级测试指标
    test_epoch_loss = test_running_loss / len(test_loader)
    test_epoch_accuracy = 100 * test_correct / test_total
    test_losses.append(test_epoch_loss)
    test_accuracies.append(test_epoch_accuracy)

    # 打印日志
    print(f'Epoch {epoch + 1}/{num_epochs}')
    print(f'Train Loss: {train_epoch_loss:.4f} | Train Acc: {train_epoch_accuracy:.2f}%')
    print(f'Test Loss: {test_epoch_loss:.4f} | Test Acc: {test_epoch_accuracy:.2f}%')
    print('-' * 30)

# 6. 可视化 Loss 和 Accuracy 曲线
plt.figure(figsize=(12, 5))

# Loss 曲线
plt.subplot(1, 2, 1)
plt.plot(train_losses, label='Train Loss')
plt.plot(test_losses, label='Test Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training vs Test Loss')
plt.legend()

# Accuracy 曲线
plt.subplot(1, 2, 2)
plt.plot(train_accuracies, label='Train Accuracy')
plt.plot(test_accuracies, label='Test Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy (%)')
plt.title('Training vs Test Accuracy')
plt.legend()

plt.show()

# 7. 可视化预测结果（选一批测试样本展示）
model.eval()
with torch.no_grad():
    # 取一个 batch 的测试数据
    images, labels = next(iter(test_loader))
    images, labels = images.to(device), labels.to(device)
    outputs = model(images)
    _, predicted = torch.max(outputs, 1)

    # 转换为 numpy（方便可视化）
    images = images.cpu().numpy()
    labels = labels.cpu().numpy()
    predicted = predicted.cpu().numpy()

    # 展示 10 个样本
    fig, axes = plt.subplots(2, 5, figsize=(15, 6))
    axes = axes.ravel()

    for i in range(10):
        # 还原归一化的图像（方便看）
        img = images[i][0] * 0.5 + 0.5
        axes[i].imshow(img, cmap='gray')
        axes[i].set_title(f'True: {labels[i]}\nPred: {predicted[i]}')
        axes[i].axis('off')

    plt.show()